from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings
from langchain_core.documents import Document
from config.config import Config
from dotenv import load_dotenv
import os

load_dotenv()

class VectorStore:
    def __init__(self, embedding_model: str = None, db_path: str = None):
        self.embedding_model = embedding_model or Config.model["EMBEDDING_MODEL"]
        self.db_path = db_path or Config.data["VECTOR_DB_PATH"]
        self.embeddings = OpenAIEmbeddings(
            model="text-embedding-3-small",
            base_url=Config.rag["base_url"],
            api_key=Config.rag["api_key"]
        )
        self.store = self._init_vector_store()
    
    def _init_vector_store(self):
        """初始化或加载向量数据库"""
        if os.path.exists(self.db_path):
            return FAISS.load_local(self.db_path, self.embeddings, allow_dangerous_deserialization=True)
        return FAISS.from_documents([], self.embeddings)
    
    def add_documents(self, documents: list[Document]):
        """添加文档到知识库"""
        self.store.add_documents(documents)
        self.store.save_local(self.db_path)
    
    def similarity_search(self, query: str, k: int = 3) -> list[str]:
        """执行相似度搜索"""
        results = self.store.similarity_search(query, k=k)
        return [doc.page_content for doc in results]

class KnowledgeManager:
    def __init__(self, vector_store: VectorStore):
        self.vector_store = vector_store
    
    def retrieve_context(self, query: str) -> list[str]:
        """检索相关知识上下文"""
        return self.vector_store.similarity_search(query)